from sys import argv 
from datetime import datetime
import os, time, subprocess, xlrd, glob
from collections import defaultdict

site = {'FZ':'福州', 'GZ':'广州', 'ZZ':'郑州', 'CS':'长沙'}
sample_type = {'BALF':'肺泡灌洗液','Blood':'血液','CSF':'脑脊液','FS':'腹水','NY':'脓液','Other':'其它','XS':'胸水'}
ss = {'BALF':'DNA','Blood':'CF','CSF':'CF','FS':'DNA','NY':'DNA','Other':'DNA','XS':'DNA'}

def RPTM(indir,job,outdir):
    #date = indir.split('/')[-1].split('_')[0]
    kkfile = glob.glob(f'{indir}/GZ.mNGS.*.mergeKK.xlsx')[0]
    ntfile = glob.glob(f'{indir}/GZ.mNGS.*.mergeNT.xlsx')[0]
    _sampletype = defaultdict(str)
    st = defaultdict(str)

    with open(f'/data/softwares/mngs_scripts/auto/autoPathoSeqPy/tag/PathoSeq_SampleType.txt') as satype:
        satype.readline()
        for l in satype:
            i = l.strip().split('\t')
            _sampletype[i[0]] = sample_type[i[1]]
            st[i[0]] = i[1]

    while 1:
        if os.path.exists(kkfile) and os.path.exists(ntfile):
            break
        time.sleep(60)

    kk = xlrd.open_workbook(kkfile)
    kksheet = kk.sheet_by_index(0)
    nt = xlrd.open_workbook(ntfile)
    ntsheet = nt.sheet_by_index(0)
    sampletype = [str(j.value).strip() for j in kksheet.row(0)]
    samplerawdata = [str(j.value).strip() for j in kksheet.row(1)]
    sampleid = [str(j.value).strip() for j in kksheet.row(2)]    
    with open(f'{indir}/historycheck.xls','w') as out:
        out.write(f'site\tjobid\truntime\tsampleid\tsampletype\tdbtype\tspecies\treads\trptm\n')
        if len(kkfile.split('/')[-2].split('_')[0]) <= 6:
            run_time = '20' + kkfile.split('/')[-2].split('_')[0]
        else:
            run_time = kkfile.split('/')[-1].split('.')[2]
        runtime = datetime.strptime(run_time, '%Y%m%d').date()
        runsite = site[kkfile.split('/')[-3]]
        for i,e in enumerate(sampleid[4:-4]): 
            if sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('CF') > -1:
                sampletype[i+4] = 'NC-CF'
            elif sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('D') > -1:
                sampletype[i+4] = 'NC-DNA'
            elif sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('R') > -1:
                sampletype[i+4] = 'NC-RNA'
            for line in range(3, kksheet.nrows):
                l = [str(j.value).strip() for j in kksheet.row(line)]                  
                if int(float(l[i+4])) > 0:
                    rptm = str('%.2f' % (int(float(l[i+4])) * 10000000/int(samplerawdata[i+4].replace(',',''))))
                    if sampleid[i+4].find('-YF') == -1:
                        if sampletype[i+4] in _sampletype:
                            out.write(f'{runsite}\t{job}\t{runtime}\t{sampleid[i+4].split(".")[0]}\t{_sampletype[sampletype[i+4]]}\tKK\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')
                            if (sampletype[i+4] == '脑脊液' and sampleid[i+4].find('CF') == -1) or sampleid[i+4].find('R') > -1:
                                continue
                            with open(f'{outdir}/{st[sampletype[i+4]]}_{ss[st[sampletype[i+4]]]}/{st[sampletype[i+4]]}_{ss[st[sampletype[i+4]]]}.batchGZ.KK.txt','a') as kkout:
                                kkout.write(f'{sampleid[i+4]}\t{l[0]}\t{l[1]}\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')    
                        else:
                            out.write(f'{runsite}\t{job}\t{runtime}\t{sampleid[i+4].split(".")[0]}\t{sampletype[i+4]}\tKK\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')
                            
        for i,e in enumerate(sampleid[4:-4]): 
            if sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('CF') > -1:
                sampletype[i+4] = 'NC-CF'
            elif sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('D') > -1:
                sampletype[i+4] = 'NC-DNA'
            elif sampleid[i+4].find('NC') > -1 and sampleid[i+4].find('R') > -1:
                sampletype[i+4] = 'NC-RNA'
            for line in range(3, ntsheet.nrows):
                l = [str(j.value).strip() for j in ntsheet.row(line)]                 
                if int(float(l[i+4])) > 0:
                    rptm = str('%.2f' % (int(float(l[i+4])) * 10000000/int(samplerawdata[i+4].replace(',',''))))
                    if sampleid[i+4].find('-YF') == -1:
                        if sampletype[i+4] in _sampletype:
                            out.write(f'{runsite}\t{job}\t{runtime}\t{sampleid[i+4].split(".")[0]}\t{_sampletype[sampletype[i+4]]}\tNT\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')
                            if (sampletype[i+4] == '脑脊液' and sampleid[i+4].find('CF') == -1) or sampleid[i+4].find('R') > -1:
                                continue
                            with open(f'{outdir}/{st[sampletype[i+4]]}_{ss[st[sampletype[i+4]]]}/{st[sampletype[i+4]]}_{ss[st[sampletype[i+4]]]}.batchGZ.NT.txt','a') as ntout:
                                ntout.write(f'{sampleid[i+4]}\t{l[0]}\t{l[1]}\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')    
                        else:
                            out.write(f'{runsite}\t{job}\t{runtime}\t{sampleid[i+4].split(".")[0]}\t{sampletype[i+4]}\tNT\t{l[2]}\t{int(float(l[i+4]))}\t{rptm}\n')
    out.close()

    script = f'''/data/mngsSYS/b/venv/bin/python /data/mngsSYS/b/app/dbCU/historycheckUpdate.py {indir}/historycheck.xls'''
    subprocess.run(script, shell=True)

if __name__ == '__main__':
    inputdir = argv[1]

    while 1:
        fn=len(glob.glob(f"{inputdir}/*.xlsx"))
        if fn==2:
            print(f"{fn} start running")
            break
        else:
            print(f"waiting ... {fn}")
            time.sleep(2)
    jobid = argv[2]
    outputdir = f'/data/POC/runmngs/autoPathoSeqWeb/batch-0713/batchGZ'
    RPTM(inputdir,jobid,outputdir) 
